import cv2
import sys
import os

import numpy as np

from face_train import Model


def get_pic(img):
    import matplotlib.pyplot as plt
    plt.imshow(img[:, :, ::-1])
    plt.show()


if __name__ == '__main__':
    if len(sys.argv) != 1:
        print("Usage:%s camera_id\r\n" % (sys.argv[0]))
        sys.exit(0)

    # 加载模型
    model = Model()
    model.load_model(file_path='./model/face.model.h5')

    # 框住人脸的矩形边框颜色
    color = (0, 0, 255)
    # 捕获指定摄像头的实时视频流
    cap = cv2.VideoCapture(0)

    # 人脸识别分类器本地存储路径
    cascade_path = "./face_model/haarcascade_frontalface_default.xml"
    # 循环检测识别人脸
    while True:
        ret, frame = cap.read()  # 读取一帧视频

        # if ret is True:
        #
        #     # 图像灰化，降低计算复杂度
        #     frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        # else:
        #     continue
        # # 使用人脸识别分类器，读入分类器
        cascade = cv2.CascadeClassifier(cascade_path)

        # 利用分类器识别出哪个区域为人脸
        faceRects = cascade.detectMultiScale(frame, scaleFactor=1.2, minNeighbors=2, minSize=(32, 32))
        if len(faceRects) > 0:
            for faceRect in faceRects:
                x, y, w, h = faceRect
                image = cv2.resize(frame, (64, 64))
                image = np.expand_dims(image, 0)
                faceID = model.predict(image)
                cv2.rectangle(frame, (x - 10, y - 10), (x + w + 10, y + h + 10), color, thickness=2)
                # face_id判断
                for i in range(len(os.listdir('./all_photo/'))):
                    if i == faceID:
                        # 文字提示是谁
                        cv2.putText(frame, os.listdir('./all_photo/')[i],
                                    (x + 30, y + 30),  # 坐标
                                    cv2.FONT_HERSHEY_SIMPLEX,  # 字体
                                    1,  # 字号
                                    (255, 0, 255),  # 颜色
                                    2)  # 字的线宽
        get_pic(frame)

        # 释放摄像头并销毁所有窗口
    cap.release()
    cv2.destroyAllWindows()
